#include <ntifs.h>
#include "ProcessEnhancer.h"
#include "utility.h"

IMPLEMENT_SINGLE_INSTANCE(ProcessEnhancer);
#define PENHANCER_SLOT_FREE 0xffffffff


ProcessEnhancer::ProcessEnhancer(void) : bProcessNotifySet(false), bOK(false), bThreadNotifySet(false)
{
	ServiceDescriptorTable* pSsdt;
	UNICODE_STRING usSsdtName;
	wchar_t* wstrSsdt = L"KeServiceDescriptorTable";
	RtlInitUnicodeString(&usSsdtName, wstrSsdt);
	for(int i = 0; i< _MAX_ENHANCE_PID_NUM; ++i)
		enhancePids[i] = PENHANCER_SLOT_FREE;
	pSsdt = (ServiceDescriptorTable*)MmGetSystemRoutineAddress(&usSsdtName);
	if(!pSsdt)
	{
		KdPrint(("dbghelper!ProcessEnhancer==>Error,Cann't get KeServiceDescriptorTable base\n"));
		return;
	}
	EnhancedServiceTable = (void**)DynamicAllocator::Instance()->Alloc(pSsdt->NumberOfServices * 4);
	KdPrint(("dbghelper!ProcessEnhancer==>SSDT base = 0x%08x, serviceNum = %d\n", pSsdt, pSsdt->NumberOfServices));
	RtlMoveMemory(EnhancedServiceTable, pSsdt->ServiceTableBase, pSsdt->NumberOfServices * 4);
	bOK = true;
	return;
}

ProcessEnhancer::~ProcessEnhancer(void)
{
	bOK = false;
	if(bProcessNotifySet)
		PsSetCreateProcessNotifyRoutine(ProcessRemoveNotify, true);
	if(bThreadNotifySet)
		PsRemoveCreateThreadNotifyRoutine(ThreadCreateNotify);
}

bool ProcessEnhancer::EnhanceProcess( UINT32 pid )
{
	PEPROCESS pEprocess;
	NTSTATUS status;
	bool bEprocessGot = false, bSlotGot = false;
	int freeId = 0;
	bool bResult = false;
	PLIST_ENTRY pThreadList;
	PLIST_ENTRY pCurrentThreadEntry;
	PETHREAD pEThread;
	if(!bOK)
		return bResult;
	do 
	{
		if(!bProcessNotifySet)
		{
			status = PsSetCreateProcessNotifyRoutine(ProcessRemoveNotify, false);
			if(!NT_SUCCESS(status))
				break;
			bProcessNotifySet = true;
		}
		if(!bThreadNotifySet)
		{
			status = PsSetCreateThreadNotifyRoutine(ThreadCreateNotify);
			if(!NT_SUCCESS(status))
				break;
			bThreadNotifySet = true;
		}
		status = PsLookupProcessByProcessId((HANDLE)pid, &pEprocess);
		if(!NT_SUCCESS(status))
			break;
		bEprocessGot = true;
		/*
		* Find Free Slot
		*/
		for(freeId = 0; freeId < _MAX_ENHANCE_PID_NUM; ++ freeId)
			if(enhancePids[freeId] == PENHANCER_SLOT_FREE)
				break;
		if(freeId == _MAX_ENHANCE_PID_NUM)
			break;
		enhancePids[freeId] = pid;
		bSlotGot = true;
		pThreadList = (PLIST_ENTRY)OffsetOf(pEprocess, 0x50);
		for(pCurrentThreadEntry = pThreadList->Flink; pCurrentThreadEntry != pThreadList; pCurrentThreadEntry = pCurrentThreadEntry->Flink)
		{
			pEThread = (PETHREAD)OffsetOf(pCurrentThreadEntry, 0-0x1b0);
			if(!EnhanceThread(pEThread))
				break;
		}
		bResult = true;
	}while(0);
	if(!bResult)
	{
		if(bSlotGot)
			enhancePids[freeId] = PENHANCER_SLOT_FREE;
	}
	if(bEprocessGot)
		ObDereferenceObject(pEprocess);
	return bResult;
}

void ProcessEnhancer::ProcessRemoveNotify( IN HANDLE ParentId,IN HANDLE ProcessId,IN BOOLEAN Create )
{
	if(Create)
		return;
	if(ProcessEnhancer::Instance()->IsProcessEnhanced((UINT32)ProcessId))
		ProcessEnhancer::Instance()->StopEnhanceProcess((UINT32)ProcessId);

}

bool ProcessEnhancer::IsProcessEnhanced( UINT32 pid )
{
	for(int id = 0; id < _MAX_ENHANCE_PID_NUM; ++id)
		if(enhancePids[id] == pid)
			return true;
	return false;
}

void ProcessEnhancer::StopEnhanceProcess( UINT32 pid )
{
	int id;
	for(id = 0; id < _MAX_ENHANCE_PID_NUM; ++id)
		if(enhancePids[id] == pid)
			break;
	if(id == _MAX_ENHANCE_PID_NUM)
		return;
	enhancePids[id] = PENHANCER_SLOT_FREE;
}

void ProcessEnhancer::ThreadCreateNotify( IN HANDLE ProcessId,IN HANDLE ThreadId,IN BOOLEAN Create )
{
	PETHREAD pEthread;
	NTSTATUS status;
	if(!Create)
		return;
	if(!ProcessEnhancer::Instance()->bOK)
		return;
	if(!ProcessEnhancer::Instance()->IsProcessEnhanced((UINT32)ProcessId))
		return;
	do 
	{
		status = PsLookupThreadByThreadId(ThreadId, &pEthread);
		if(!NT_SUCCESS(status))
			break;
		ProcessEnhancer::Instance()->EnhanceThread(pEthread);
		ObDereferenceObject(pEthread);
	} while (0);
}

bool ProcessEnhancer::EnhanceThread( PETHREAD pEthread )
{
	if(!bOK)
		return false;
	ServiceDescriptorTable* pThreadSSDT = (ServiceDescriptorTable*)FetchDataDWORD(pEthread, 0xe0);
	pThreadSSDT->ServiceTableBase = ProcessEnhancer::Instance()->EnhancedServiceTable;
	KdPrint(("dbghelper!ProcessEnhancer::EnhanceThread==>change thread service table ok,ethread=0x%08x\n", pEthread));
	return true;
}